import os
import json
import cv2

# 定义目录路径
base_dir = "288230376152431921"
annotations_dir = os.path.join(base_dir, "annotations")
images_dir = os.path.join(base_dir, "images")
output_dir = os.path.join(base_dir, "outputs")

# 创建输出文件夹（如果不存在）
os.makedirs(output_dir, exist_ok=True)

# 遍历所有 JSON 标注文件
for filename in os.listdir(annotations_dir):
    if not filename.endswith(".json"):
        continue

    json_path = os.path.join(annotations_dir, filename)

    # 读取 JSON 标注文件
    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    # 构造对应图像文件路径
    image_name = os.path.splitext(filename)[0] + ".jpg"  # 可根据实际图像格式修改
    image_path = os.path.join(images_dir, image_name)

    if not os.path.exists(image_path):
        print(f"图像不存在：{image_path}")
        continue

    # 加载图像
    image = cv2.imread(image_path)

    # 遍历所有标注框
    for shape in data.get("shapes", []):
        if shape.get("shape_type") != "rectangle":
            continue

        points = shape.get("points", [])
        if len(points) != 2:
            continue

        # 获取左上角和右下角坐标
        (x1, y1), (x2, y2) = points
        pt1 = (int(x1), int(y1))
        pt2 = (int(x2), int(y2))

        # 绘制矩形框
        cv2.rectangle(image, pt1, pt2, color=(0, 255, 0), thickness=2)

        # 可选：绘制 label 名称
        label = shape.get("label", "")
        cv2.putText(image, label, (pt1[0], pt1[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 
                    0.5, (0, 255, 0), 1, lineType=cv2.LINE_AA)

    # 保存输出图像
    output_path = os.path.join(output_dir, image_name)
    cv2.imwrite(output_path, image)
    print(f"处理完成：{output_path}")
